import fitz

def get_dirs(bookmark_txt_path, offset):
    file_obj = open(bookmark_txt_path, encoding="utf-8")
    all_lines = file_obj.readlines()
    # print(all_lines)
    dirs = []
    for line in all_lines:
        item = []
        lvl = 0
        content = ""
        idx = 0
        flag1 = True
        flag2 = False
        flag3 = False

        for c in line:
            # 计算目录级别
            if c == ' ' and flag1:
                lvl = lvl + 1
            elif flag1:
                # print("haha")
                flag1 = False
                flag2 = True
                # break
            # 计算目录内容
            if (c != '\t' and c != ' ' and c != '\n') and flag2:
                content += c
            elif flag2:
                # print("lala")
                flag2 = False
                flag3 = True
                # break;
            # 计算页码
            if (c >= '0' and c <= '9') and flag3:
                idx = idx * 10 + int(c)
            # else:
                # break;

        print(content + ' ' + str(idx + offset))
        item.append((int)(lvl / 4) + 1)
        item.append(content)
        item.append(idx + offset)
        dirs.append(item)

    file_obj.close()
    return dirs

def addBookmark(pdf_path, bookmark_txt_path, offset):
    # 读取 pdf 文件
    doc = fitz.open(pdf_path)
    # 获得目录
    dirs = get_dirs(bookmark_txt_path, offset)

    toc = []
    # 生成目录
    for item in dirs:
        toc.append(item)
    # 替换目录
    doc.setToC(toc)
    # 保存新文件
    doc.save('H://Git//SundriesRoom//电子书目录生成//矩阵分析.pdf')

if __name__ == "__main__":
    addBookmark("H://Git//SundriesRoom//电子书目录生成//test.pdf", "H://Git//SundriesRoom//电子书目录生成//toc.txt", 11)